import numpy as np

class FilterTarget:
    """
    筛选出子任务所需的指定类型的对象数据
    需要与CropTarget配合使用
    args:
        keys: 指定了从data中读取数据的关键字和保存数据的关键字
            格式: keys={"in":"your_in_keyword",
                        "out":"your_save_keyword"
                        }
    """

    def __init__(self, keys, target_cls):
        self.keys = keys
        self.target_cls = target_cls

    def __call__(self, data: dict):
        res = []
        for crop_data in data[self.keys["in"]]:
            filter_data = {
                "imgs": [],
                "det_ids": [],
                "clses": [],
                "boxes": [],
                # "org_boxes":[],
            }
            for idx_, [crop_input_, crop_idx_, crop_cls_, expend_box_] in enumerate(
                zip(
                    crop_data["imgs"],
                    crop_data["det_ids"],
                    crop_data["clses"],
                    crop_data["boxes"],
                )
            ):

                if crop_cls_ in self.target_cls:
                    filter_data["imgs"].append(crop_input_)
                    filter_data["det_ids"].append(crop_idx_)
                    filter_data["clses"].append(crop_cls_)
                    filter_data["boxes"].append(expend_box_)
            res.append(filter_data)

        data[self.keys["out"]] = res

        return data

class MMFilterTarget:
    """
    筛选出子任务所需的指定类型的对象数据
    """

    def __init__(self, keys, target_cls):
        self.keys = keys
        self.target_cls = target_cls

    def __call__(self, data: dict):
        
        batch_bboxes = data[self.keys["in_det"]]["boxes"]
        target_bboxes = []
        target_masks=[]
        for bboxes in batch_bboxes:
            if len(bboxes)<=0:
                target_bboxes.append(np.empty_like(bboxes))
                target_masks.append([])
                continue
            mask = np.isin(bboxes[:, 0], self.target_cls)
            res = bboxes[mask]
            target_bboxes.append(res)
            target_masks.append(mask)
        data[self.keys["out_det"]] = {
            "boxes": target_bboxes,
            "masks": target_masks
        }

        return data
